import requests
import tempfile
import os
from markitdown import MarkItDown


def download_and_parse_arxiv_pdf(url):
    """
    下载 arXiv PDF 并使用 MarkItDown 解析内容

    Args:
        url: arXiv PDF 链接，例如 'https://arxiv.org/pdf/2510.09060'

    Returns:
        dict: 包含 parsed_content, success, error 的字典
    """
    result = {
        'parsed_content': None,
        'success': False,
        'error': None
    }

    try:
        # 1. 下载 PDF
        print(f"正在下载: {url}")
        response = requests.get(url, timeout=30)
        response.raise_for_status()

        content = response.content
        print(f"下载完成，文件大小: {len(content)} bytes")

        # 2. 使用 MarkItDown 解析
        md = MarkItDown()

        # 创建临时文件
        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as tmp_file:
            tmp_file.write(content)
            tmp_file.flush()
            tmp_path = tmp_file.name

        try:
            # 解析 PDF
            print("正在解析 PDF...")
            parse_result = md.convert(tmp_path)
            result['parsed_content'] = parse_result.text_content
            result['success'] = True
            print("解析成功！")

        except Exception as parse_error:
            result['error'] = f"解析失败: {str(parse_error)}"
            print(result['error'])

        finally:
            # 清理临时文件
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    except requests.exceptions.RequestException as e:
        result['error'] = f"下载失败: {str(e)}"
        print(result['error'])
    except Exception as e:
        result['error'] = f"未知错误: {str(e)}"
        print(result['error'])

    return result


# 使用示例
if __name__ == "__main__":
    # arXiv PDF 链接
    pdf_url = "https://arxiv.org/pdf/2510.09060"

    # 下载并解析
    result = download_and_parse_arxiv_pdf(pdf_url)

    if result['success']:
        print("\n" + "=" * 50)
        print("解析内容预览（前500字符）:")
        print("=" * 50)
        print(result['parsed_content'][:500])
        print("...")

        # 可选：保存到文件
        output_file = "arxiv_paper.md"
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(result['parsed_content'])
        print(f"\n完整内容已保存到: {output_file}")
    else:
        print(f"\n处理失败: {result['error']}")